# app/core/db.py
import asyncio
import logging
from contextlib import asynccontextmanager
from typing import Annotated, AsyncGenerator

from fastapi import Depends
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker

from app.core import settings

logger = logging.getLogger("async_db")

_async_engine: AsyncEngine | None = None


def init_async_engine() -> AsyncEngine:
    global _async_engine
    if _async_engine is not None:
        return _async_engine

    _async_engine = create_async_engine(
        str(settings.SQLALCHEMY_ASYNC_DATABASE_URI),
        pool_size=settings.GCE_DB.POOL_SIZE or 20,
        max_overflow=settings.GCE_DB.MAX_OVERFLOW or 10,
        pool_timeout=settings.GCE_DB.POOL_TIMEOUT or 30,
        pool_recycle=settings.GCE_DB.POOL_RECYCLE or 1800,
        pool_pre_ping=True,
        echo=settings.GCE_DB.ECHO,
        future=True,
        connect_args=(
            {"command_timeout": 60, "server_settings": {"application_name": settings.PROJECT_NAME}}
            if settings.GCE_DB.TYPE == "postgres"
            else {}
        ),
    )
    logger.info("数据库引擎初始化成功")
    return _async_engine


def get_async_engine() -> AsyncEngine:
    if _async_engine is None:
        raise RuntimeError("数据库引擎未初始化")
    return _async_engine


def get_session_factory() -> sessionmaker[AsyncSession]:
    return sessionmaker(
        bind=get_async_engine(),
        class_=AsyncSession,
        autocommit=False,
        autoflush=False,
        expire_on_commit=False,
    )


async def get_db_session() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI 依赖注入：手动控制事务"""
    session = get_session_factory()()
    try:
        yield session
    except Exception as e:
        logger.error(f"数据库会话异常: {e}")
        await session.rollback()
        raise
    finally:
        await session.close()


SessionDep = Annotated[AsyncSession, Depends(get_db_session)]


@asynccontextmanager
async def auto_commit_session() -> AsyncGenerator[AsyncSession, None]:
    """后台任务使用：自动提交事务"""
    session = get_session_factory()()
    try:
        yield session
        await session.commit()
    except Exception:
        await session.rollback()
        raise
    finally:
        await session.close()


async def check_db_connection(timeout: float = 5.0) -> bool:
    try:
        async with get_async_engine().connect() as conn:
            await asyncio.wait_for(conn.execute(text("SELECT 1")), timeout=timeout)
        return True
    except (asyncio.TimeoutError, Exception) as e:
        logger.error(f"数据库连接检查失败: {e}")
        return False


async def close_engine() -> None:
    global _async_engine
    if _async_engine:
        await _async_engine.dispose()
        _async_engine = None
        logger.info("数据库引擎已关闭")
